import pathway as pw

from application.db.mysql_db.nsfc.NsfcInfoTypeDict import NsfcInfoTypeDict


class InfoProcessPipeline:
    apply_code_dict = {
        "数理科学部": "A",
        "数学物理科学部": "A",
        "化学科学部": "B",
        "生命科学部": "C",
        "地球科学部": "D",
        "工程与材料科学部": "E",
        "信息科学部": "F",
        "管理科学部": "G",
        "医学科学部": "H",
    }
    info_type_dict = {item.info_type_name: item.info_type_id for item in
                      NsfcInfoTypeDict.select(NsfcInfoTypeDict.info_type_id, NsfcInfoTypeDict.info_type_name)}

    def apply(self, value):
        # 添加 info_type_id 列
        result = value.with_columns(
            info_type_id=pw.apply(self.info_type_id, value.info_tags)
        )

        # 添加 apply_code 列
        result = result.with_columns(
            apply_code=pw.apply(self.apply_code, result.info_tags)
        )
        result = result.with_columns(
            is_del=0
        )

        return result

    def info_type_id(self, info_tags):
        if not info_tags:
            return 0
        info_tags = info_tags._value
        tag_value = info_tags['tag_value']
        if not tag_value:
            return 0
        info_type_id = self.info_type_dict.get(tag_value[0], 0)
        if info_type_id:
            return info_type_id
        return 0

    def apply_code(self, info_tags):
        if not info_tags:
            return "*"
        info_tags = info_tags._value
        tag_value = info_tags['tag_value']
        if not tag_value:
            return "*"
        if len(tag_value) < 2:
            return "*"
        return self.apply_code_dict.get(tag_value[1], "*")


if __name__ == '__main__':
    data = NsfcInfoTypeDict.select(NsfcInfoTypeDict.info_type_id)
    for i in data:
        print(i.info_type_id)
